import math
from paddle.fluid.dygraph.learning_rate_scheduler import LearningRateDecay

__all__ = ["CosineLR"]


class CosineLR(LearningRateDecay):

    def __init__(self,
                 learning_rate,
                 step_each_epoch,
                 epochs,
                 min_lr,
                 warmup_min_lr,
                 warmup_epoch=5,
                 begin=1,
                 step=1,
                 dtype='float32'):
        super().__init__(begin, step, dtype)
        self.learning_rate = learning_rate
        self.step_each_epoch = step_each_epoch
        self.epochs = epochs
        self.min_lr = min_lr
        self.warmup_min_lr = warmup_min_lr
        self.warmup_epoch = warmup_epoch

    def step(self):
        cur_epoch = self.step_num / self.step_each_epoch
        if cur_epoch <= self.warmup_epoch:
            max_lr = self.learning_rate
            min_lr = self.warmup_min_lr
            lr = (max_lr - min_lr) * cur_epoch / self.warmup_epoch + min_lr
        elif cur_epoch <= self.epochs:
            cur_epoch = cur_epoch - self.warmup_epoch
            epochs = self.epochs - self.warmup_epoch
            max_lr = self.learning_rate
            min_lr = self.min_lr
            lr = (max_lr - min_lr) * 0.5 * (
                    math.cos(cur_epoch * math.pi / epochs) + 1) + min_lr
        else:
            lr = self.min_lr
        return lr
